import os.path

import gradio as gr

try:
    from conf.config import BASE_DIR, logger
    from utils import chat_building
except ModuleNotFoundError:
    import os
    import sys
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # 离开IDE也能正常导入自己定义的包
    from conf.config import BASE_DIR, logger
    from utils import chat_building


def get_answer(question: str, top_k: int, chain_type: str):
    """
    获取答复
    :param question: 待答复问题
    :param top_k: 要查找前几个相关文档
    :param chain_type: 要使用的langchain问答链的类型
    :return:
    """
    answer, answer_details = chat_building.get_answer(question=question, top_k=top_k, chain_type=chain_type)
    answer_details = answer_details["relevant_docs"]

    logger.info(f"问题：{question}")
    logger.info(f"top_k：{top_k}")
    logger.info(f"chain_type：{chain_type}")
    logger.info(f"答复：\n{answer}")
    logger.info(f"详情：\n{answer_details}")

    return answer.replace("SOURCES", "文本块编号"), answer_details


def main():
    with gr.Blocks(theme=gr.themes.Base(), title="Chat Building") as app:
        gr.Markdown("# <center>Chat Building")

        # 输入组件
        ele_text_box_question = gr.Textbox(label="你想问什么", placeholder="请输入你的问题", lines=2)

        with gr.Group():
            ele_slider_tok_k = gr.Slider(1, 6, value=4, step=1,
                                         label="要检索前几的相关文本进行答复", info="在1到6间选择")
            ele_radio_chain_type = gr.Radio(["stuff", "map_reduce"],
                                            value="stuff",
                                            label="与GPT的交互方式", info="stuff: 直接填充；map_reduce: 先摘要再填充")

        # 输入示例
        gr.Examples(
            label="输入示例，点击可快速填充",
            examples=[
                ["万科科技之光的咨询电话是多少？", 4, "stuff"],
                ["万科科技之光的容积率是多少？", 6, "stuff"]
            ],
            inputs=[ele_text_box_question, ele_slider_tok_k, ele_radio_chain_type]
        )

        # 动作按钮
        ele_button = gr.Button("提交", variant="primary")

        # 输出组件
        ele_text_box_answer = gr.Textbox(label="答复（仅供参考）", lines=1)
        ele_dataframe_answer_details = gr.Dataframe(
            label="答复的文件源",
            headers=["文本块编号", "向量距离", "文本内容", "文件源"],
            datatype=["str", "str", "str", "str"],
            row_count=(1, "dynamic"),
            col_count=(4, "fixed"),
            wrap=True
        )

        # 事件绑定
        ele_button.click(fn=get_answer,
                         inputs=[ele_text_box_question, ele_slider_tok_k, ele_radio_chain_type, ],
                         outputs=[ele_text_box_answer, ele_dataframe_answer_details, ])

        app.queue(concurrency_count=10)

        app.launch(show_api=False, server_name="0.0.0.0", server_port=9077,
                   favicon_path=os.path.join(BASE_DIR, "logo.png"))


if __name__ == '__main__':
    main()
